
import pymysql
from sshtunnel import SSHTunnelForwarder
from typing import List, Tuple, Optional


class MySQLSSHClient:
    def __init__(
        self,
        ssh_host: str ="8.211.48.200",
        ssh_port: int = 22,
        ssh_user: str = "cys",
        ssh_password: str = "Keepgo123@cys",
        ssh_pkey: Optional[str] = None,
        mysql_host: str = "rm-gw8a4drz5t1dde7g3.mysql.germany.rds.aliyuncs.com" ,
        mysql_port: int = 3306,
        mysql_user: str = "root",
        mysql_password: str = "suKeepgo123@cys",
        mysql_db: str = "cfx_t",
        local_bind_port: int = 3307
    ):
        """
        初始化MySQL SSH客户端
        
        :param ssh_host: SSH服务器地址
        :param ssh_port: SSH端口
        :param ssh_user: SSH用户名
        :param ssh_password: SSH密码（与ssh_pkey二选一）
        :param ssh_pkey: SSH私钥路径（与ssh_password二选一）
        :param mysql_host: MySQL主机地址（SSH服务器视角）
        :param mysql_port: MySQL端口
        :param mysql_user: MySQL用户名
        :param mysql_password: MySQL密码
        :param mysql_db: 数据库名
        :param local_bind_port: 本地绑定端口
        """
        self.ssh_host = ssh_host
        self.ssh_port = ssh_port
        self.ssh_user = ssh_user
        self.ssh_password = ssh_password
        self.ssh_pkey = ssh_pkey
        self.mysql_host = mysql_host
        self.mysql_port = mysql_port
        self.mysql_user = mysql_user
        self.mysql_password = mysql_password
        self.mysql_db = mysql_db
        self.local_bind_port = local_bind_port
        
        # 连接对象
        self.ssh_tunnel: Optional[SSHTunnelForwarder] = None
        self.db_conn: Optional[pymysql.connections.Connection] = None

    def connect(self) -> bool:
        """建立SSH隧道和MySQL连接"""
        try:
            # 创建SSH隧道
            self.ssh_tunnel = SSHTunnelForwarder(
                (self.ssh_host, self.ssh_port),
                ssh_username=self.ssh_user,
                ssh_password=self.ssh_password,
                ssh_pkey=self.ssh_pkey,
                remote_bind_address=(self.mysql_host, self.mysql_port),
                local_bind_address=("127.0.0.1", self.local_bind_port)
            )
            self.ssh_tunnel.start()
            #print(f"SSH隧道已启动，本地端口: {self.ssh_tunnel.local_bind_port}")

            # 连接MySQL
            self.db_conn = pymysql.connect(
                host="127.0.0.1",
                port=self.ssh_tunnel.local_bind_port,
                user=self.mysql_user,
                password=self.mysql_password,
                database=self.mysql_db,
                charset="utf8mb4"
            )
            #print("MySQL连接成功")
            return True
        except Exception as e:
            print(f"连接失败: {str(e)}")
            self.close()  # 确保资源释放
            return False

    def query(self, sql: str) -> Tuple[List[Tuple], List[str]]:
        """
        执行查询SQL
        
        :param sql: SQL查询语句
        :return: 元组(查询结果数据, 列名列表)
        """
        if not self.db_conn or not self.ssh_tunnel:
            raise Exception("请先调用connect()建立连接")

        try:
            with self.db_conn.cursor() as cursor:
                cursor.execute(sql)
                # 获取查询结果
                data = cursor.fetchall()
                # 获取列名
                columns = [desc[0] for desc in cursor.description]
                return data, columns
        except pymysql.MySQLError as e:
            print(f"查询失败: {str(e)}")
            raise  # 抛出异常让调用者处理

    def close(self) -> None:
        """关闭所有连接"""
        # 关闭MySQL连接
        if self.db_conn and self.db_conn.open:
            self.db_conn.close()
            #print("MySQL连接已关闭")
        
        # 关闭SSH隧道
        if self.ssh_tunnel:
            self.ssh_tunnel.stop()
            #print("SSH隧道已关闭")

    def __enter__(self):
        """支持上下文管理器"""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器退出时关闭连接"""
        self.close()


# 使用示例
if __name__ == "__main__":
    

    # 方法1: 常规使用
    client = MySQLSSHClient()
    if client.connect():
        try:
            data, columns = client.query("SELECT * FROM user_channel_item LIMIT 10;")
            print("列名:", columns)
            print("数据:")
            for row in data:
                print(row)
        finally:
            client.close()


